package main

import (
	"context"
	"fmt"
	"github.com/google/uuid"
	"go.mongodb.org/mongo-driver/bson"
	"log"
	"sync"
	"time"

	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// TaskStatus 任务状态
type TaskStatus int

// 定义任务状态常量
const (
	Running TaskStatus = iota
	Failed
	Terminated
	ManualTerminated
	ManualTerminatedSuccess
	Success
)

// Task 任务结构体
type Task struct {
	ID         string     `json:"task_uuid" bson:"task_uuid"`
	CreateTime time.Time  `json:"create_time" bson:"create_time"`
	UpdateTime time.Time  `json:"update_time" bson:"update_time"`
	Status     TaskStatus `json:"status" bson:"status"`
	Progress   int        `json:"progress" bson:"progress"`
	Params     string     `json:"params" bson:"params"`
	RetryCount int        `json:"retry_count" bson:"retry_count"`
}

// TaskScheduler 任务调度器
type TaskScheduler struct {
	mutex       sync.Mutex    // 互斥锁
	mongoClient *mongo.Client // MongoDB 客户端
}

// NewTaskScheduler 创建新的任务调度器
func NewTaskScheduler(mongoURI string) (*TaskScheduler, error) {
	// 连接 MongoDB 数据库
	client, err := mongo.Connect(context.Background(), options.Client().ApplyURI(mongoURI))
	if err != nil {
		return nil, err
	}

	return &TaskScheduler{
		mongoClient: client,
	}, nil
}

// AddTask 添加任务到调度器
func (ts *TaskScheduler) AddTask(params string) string {
	ts.mutex.Lock()
	defer ts.mutex.Unlock()

	// 生成任务ID
	taskID := uuid.New().String()

	// 创建新任务
	task := &Task{
		ID:         taskID,
		CreateTime: time.Now(),
		UpdateTime: time.Now(),
		Status:     Running,
		Progress:   0,
		Params:     params,
	}

	err := ts.AddTaskToMongo(task)
	if err != nil {
		log.Printf("新增任务插入mongo失败")
	}
	// 异步执行任务
	go ts.runTask(task)

	time.Sleep(time.Second * 1000)
	return taskID
}

// runTask 异步执行任务
func (ts *TaskScheduler) runTask(task *Task) {
	task.UpdateTime = time.Now()
	err := ts.UpdateTask(task)
	if err != nil {
		log.Printf("Failed to update task info in MongoDB: %v\n", err)
	}

	// 模拟任务执行过程，这里可以根据实际需求替换为真实的任务执行逻辑
	for i := 1; i <= 100; i++ {
		time.Sleep(1000 * time.Millisecond)
		task.Progress = i
		task.UpdateTime = time.Now()
		err = ts.UpdateTask(task)
		if err != nil {
			log.Printf("Failed to update task info in MongoDB: %v\n", err)
			return
		}
	}
	// 判断任务运行情况来设置任务状态
	// 模拟任务执行结果，这里假设任务执行成功
	err = ts.storeTaskResult(task)
	if err != nil {
		log.Printf("Failed to write  task result in MongoDB: %v\n", err)
		task.Status = Failed
		task.UpdateTime = time.Now()
		err = ts.UpdateTask(task)
		if err != nil {
			log.Printf("Failed to update task info in MongoDB: %v\n", err)
			return
		}
		return
	}

	task.Status = Success
	task.UpdateTime = time.Now()
	err = ts.UpdateTask(task)
	if err != nil {
		log.Printf("Failed to update task info in MongoDB: %v\n", err)
		return
	}
	log.Printf("%#v", task)
}

// GetTaskStatus 获取任务状态
func (ts *TaskScheduler) GetTaskStatus(taskID string) (TaskStatus, error) {
	ts.mutex.Lock()
	defer ts.mutex.Unlock()

	task, err := ts.getTaskFromMongo(taskID)
	if err != nil {
		return 0, fmt.Errorf("task with ID %s not found", taskID)
	}

	return task.Status, nil
}

// RetryTask 重试任务
func (ts *TaskScheduler) RetryTask(taskID string) error {
	ts.mutex.Lock()
	defer ts.mutex.Unlock()

	task, err := ts.getTaskFromMongo(taskID)
	if err != nil {
		return fmt.Errorf("task with ID %s not found", taskID)
	}

	// 检查任务状态是否为 Failed，只有失败的任务才能重试
	if task.Status != Failed {
		return fmt.Errorf("task with ID %s is not in failed status", taskID)
	}

	// 最多重试 3 次
	if task.RetryCount >= 3 {
		// 超过最大重试次数，将任务状态设置为 Terminated
		task.Status = Terminated
		task.UpdateTime = time.Now()
		err := ts.UpdateTask(&task)
		if err != nil {
			log.Printf("Failed to update task info in MongoDB: %v\n", err)
		}
		return fmt.Errorf("task with ID %s has reached maximum retry count", taskID)
	}

	// 重试间隔 5 分钟
	time.Sleep(5 * time.Minute)

	// 重置任务状态和进度
	task.Status = Running
	task.Progress = 0
	task.RetryCount++

	task.UpdateTime = time.Now()
	err = ts.UpdateTask(&task)
	if err != nil {
		log.Printf("Failed to update task info in MongoDB: %v\n", err)
	}

	// 异步执行任务
	go ts.runTask(&task)

	return nil
}

type TaskResult struct {
	TaskID     string      `json:"task_uuid" bson:"task_uuid"`
	FinishTime time.Time   `json:"finish_time" bson:"finish_time"`
	ResultInfo interface{} `json:"result_info" bson:"result_info"`
}

// getTaskFromMongo 从 MongoDB 中获取任务信息
func (ts *TaskScheduler) getTaskFromMongo(taskID string) (task Task, err error) {

	// 从 MongoDB 中获取任务信息
	collection := ts.mongoClient.Database("taskdb").Collection("task_info")
	filter := bson.M{"task_uuid": taskID}
	err = collection.FindOne(context.Background(), filter).Decode(&task)
	if err != nil {
		log.Printf("Failed to get task from MongoDB: %v\n", err)
		return
	}
	return task, err
}

// UpdateTask 更新任务信息到 MongoDB
func (ts *TaskScheduler) UpdateTask(task *Task) error {
	collection := ts.mongoClient.Database("taskdb").Collection("task_info")

	filter := bson.M{"task_uuid": task.ID}
	update := bson.M{"$set": task}

	_, err := collection.UpdateOne(context.Background(), filter, update)
	if err != nil {
		return fmt.Errorf("failed to update task in MongoDB: %v", err)
	}

	return nil
}

// AddTaskToMongo 将任务添加到 MongoDB 中
func (ts *TaskScheduler) AddTaskToMongo(task *Task) error {
	// 将任务添加到 MongoDB 中
	collection := ts.mongoClient.Database("taskdb").Collection("task_info")
	_, err := collection.InsertOne(context.Background(), task)
	if err != nil {
		return fmt.Errorf("failed to insert task in MongoDB: %v", err)
	}
	return nil
}

// storeTaskResult 将任务结果存入 MongoDB
func (ts *TaskScheduler) storeTaskResult(task *Task) error {
	collection := ts.mongoClient.Database("taskdb").Collection("task_result")

	result := &TaskResult{
		TaskID:     task.ID,
		FinishTime: time.Now(),
		ResultInfo: "This is a mock result.s",
	}

	_, err := collection.InsertOne(context.Background(), result)
	if err != nil {
		log.Printf("Failed to store task result in MongoDB: %v\n", err)
		return err
	}
	return nil
}
